/*
 * Copyright (c) 2005, 2014 vacoor
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 */
package org.vacoor.nothing.wdb.websocket;

import com.google.common.collect.Maps;
import com.pty4j.PtyProcess;
import com.pty4j.WinSize;
import com.pty4j.util.PtyUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.vacoor.nothing.common.codec.Hex;
import org.vacoor.nothing.common.misc.util.Platform;
import org.vacoor.nothing.common.util.IOUtils;
import org.vacoor.nothing.common.util.Throwables;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.*;
import java.net.URL;
import java.nio.charset.Charset;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Map;

/**
 * Terminal WebSocket Endpoint
 *
 * @author vacoor
 */
@ServerEndpoint("/terminal/{id}")
public class TerminalEndpoint {
    private static final Logger LOG = LoggerFactory.getLogger(TerminalEndpoint.class);
    private static final Charset DEFAULT_CHARSET = Charset.forName("UTF-8");
    private static final String MOTD = "" +
            " __    __    _ _               _____                    _             _ \r\n" +
            "/ / /\\ \\ \\__| | |__           /__   \\___ _ __ _ __ ___ (_)_ __   __ _| |\r\n" +
            "\\ \\/  \\/ / _` | '_ \\   _____    / /\\/ _ \\ '__| '_ ` _ \\| | '_ \\ / _` | |\r\n" +
            " \\  /\\  / (_| | |_) | |_____|  / / |  __/ |  | | | | | | | | | | (_| | |\r\n" +
            "  \\/  \\/ \\__,_|_.__/           \\/   \\___|_|  |_| |_| |_|_|_| |_|\\__,_|_|\r\n" +
            "\r\n";

    private PtyProcess pty;

    @OnOpen
    public void onOpen(@PathParam("id") String id, final Session session) throws IOException {
        LOG.debug("Session {} open", session.getId());

        final RemoteEndpoint.Basic re = session.getBasicRemote();

        Map<String, String> envs = Maps.newHashMap(System.getenv());
        envs.put("TERM", "xterm-color");
        envs.put("LANG", "en_US.UTF-8");
        envs.put("LC_CTYPE", "UTF-8");  // must be, shell charset {@see} man bash
        envs.put("BASH_SOURCE", "/etc/bashrc:~/.bashprofile:~/.bashrc");
        envs.put("SHLVL", "1");
        envs.put("TERM_PROGRAM", "Wdb_Terminal");

        String[] command = Platform.isWindows() ? new String[]{"cmd.exe"} : new String[]{"/bin/bash", "--login"};

        pty = PtyProcess.exec(command, envs, System.getProperty("user.home"));
        pty.setWinSize(new WinSize(80, 24));
        re.sendText(MOTD);

        PtyStarter.of(pty, session).start();
        session.setMaxIdleTimeout(10 * 60 * 1000);
    }

    @OnMessage
    public void onMessage(String message, Session session) throws IOException {
        LOG.debug("received message: {} from session {}", message, session.getId());

        if (pty != null && pty.isRunning()) {
            if (message.startsWith("resize:")) {
                message = message.substring(7);
                String[] size = message.split(",");
                try {
                    WinSize winSize = new WinSize(Integer.valueOf(size[0]), Integer.valueOf(size[1]));
                    pty.setWinSize(winSize);
                } catch (NumberFormatException ignore) {
                    // ignore
                }
            } else {
                OutputStream out = pty.getOutputStream();
                out.write(message.getBytes(DEFAULT_CHARSET));
                out.flush();
            }
        }
    }

    @OnClose
    public void onClose(Session session, CloseReason closeReason) {
        LOG.debug("Session {} closed because of {}", session.getId(), closeReason);

        if (pty != null && pty.isRunning()) {
            pty.destroy();
            pty = null;
        }
    }

    @OnError
    public void onError(Throwable ex, Session session) {
        LOG.warn("Session exception: {}", ex);
        if (session.isOpen()) {
            try {
                session.close();
            } catch (IOException e) {
                // ignore
            }
        }
    }

    /* *****************************
     *
     * *****************************/
    private static class PtyStarter extends Thread {
        private final PtyProcess pty;
        private final Session session;

        private static PtyStarter of(PtyProcess pty, Session session) {
            return new PtyStarter(pty, session);
        }

        private PtyStarter(PtyProcess pty, Session session) {
            this.pty = pty;
            this.session = session;
        }

        @Override
        public void run() {
            try {
                InputStream is = pty.getInputStream();
                RemoteEndpoint.Basic re = session.getBasicRemote();
                int len;
                byte[] buf = new byte[1024];
                while (-1 < (len = is.read(buf))) {
                    re.sendText(new String(buf, 0, len, DEFAULT_CHARSET));
                }

                /*- logout/exit --> finally */
            } catch (IOException ignore) {
                // ignore
            } finally {
                if (session.isOpen()) {
                    try {
                        session.close(new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "Logout"));
                    } catch (IOException ignore) {
                        // ignore
                    }
                }
            }
        }
    }

    static {
        try {
            extractPtyLibrary();
        } catch (Exception e) {
            Throwables.unchecked(e);
        }
    }

    /**
     * libpty 默认从 jar 所在文件夹加载 library, 这里修改下方便打包到 jar 里面
     */
    private static void extractPtyLibrary() throws Exception {
        final String ptyLibFolder = PtyUtil.getPtyLibFolderPath();  // 获取 PtyLib 基准文件夹
        final String libRelativePath = PtyUtil.resolveNativeLibrary(new File("libpty")).getPath();  // 获取相对基准文件夹的路径
        final File libFile = new File(ptyLibFolder, libRelativePath);
        URL url = TerminalEndpoint.class.getClassLoader().getResource(libRelativePath);

        if (libFile.exists() && md5(libFile).equals(url.openStream())) {
            LOG.debug("libpty library {} is already exists.", libFile);
            return;
        }

        File libFolderFile = libFile.getParentFile();
        if (!libFolderFile.exists()) {
            libFolderFile.mkdirs();
        }
        libFile.deleteOnExit();

        LOG.debug("extra libpty library {} to {}", libRelativePath, libFile);
        FileOutputStream libOutStream = new FileOutputStream(libFile);

        IOUtils.flow(url.openStream(), libOutStream, true, true);
    }


    /* *****************************
     *      HELPER METHOD
     * *****************************/
    private static String md5(File file) throws IOException {
        return null != file && file.exists() ? md5(new FileInputStream(file)) : null;
    }

    private static String md5(InputStream in) throws IOException {
        try {
            MessageDigest md5 = MessageDigest.getInstance("MD5");

            try {
                byte[] buf = new byte[8192];
                int len;
                while ((len = in.read(buf)) >= 0) {
                    md5.update(buf, 0, len);
                }
                return Hex.encode(md5.digest());

            } finally {
                in.close();
            }
        } catch (NoSuchAlgorithmException e) {
            throw new AssertionError(e);
        }
    }
}
